#/usr/bin/python
# -*- coding: utf-8 -*-

from evtk import hl, vtk
import numpy as np
import parameters
import re

class Parser(object):
    
    def __init__(self, domain, filename, paramfile=None):
        
        self.domain = domain
        self.filename = filename
        self.paramFile = paramfile
        
        self.node_fields = ['dof', 'force', 'residual']
        self.elem_fields = ['stress', 'strain']
        
    def writeOutputFile(self):
        '''
            write the saved fields into a file whose format is defined as parameter.
            Fields are saved as Field objects associated to a mesh. The Field objects are associating state identifiers
            with fields on mesh. The mapping of states (type "STEP" or "MODE", ordering) is hold by the domain object. 
        '''
        return 0
    
class ParserVTK(Parser):
    
    cells_type_dict = {
    'S3': vtk.VtkTriangle.tid,
    'S4': vtk.VtkQuad.tid,
    'S6': vtk.VtkQuadraticTriangle.tid,
    'S8': vtk.VtkQuadraticQuad.tid,
    'B2': vtk.VtkLine.tid,
    'B3': vtk.VtkQuadraticEdge.tid,
    'R2': vtk.VtkLine.tid,          
    }                  
    
    cells_connec_dict = {
    'S3': [0,1,2],
    'S4': [0,1,2,3],
    'S6': [0,1,2,3,4,5],
    'S8': [0,1,2,3,4,5,6,7],
    'B2': [0,1],
    'B3': [0,1,2],
    'R2': [0,1], 
    }
    
    def writeOutputFile(self, levelfilter=[]):
        
        param = parameters.load()
        
        print param.OUTPUT_FIELDS
        
        # for each domain:
        try:
            for domain in self.domain.domains: # this is a global domain
                parser = ParserVTK(domain, self.filename + domain.UId)
                parser.writeOutputFile()
        except:
            pass # non global domain
        
        # states by type:
        for type in self.domain.statesMap.keys():
            
            statekeys = self.domain.statesMap[type]
            if not len(statekeys): continue
        
            if type == 'STEP':
                iterpattern = re.compile(r"ST(\d+\.*\d*)-IN(\d+)-IT(\d+)")
                incrpattern = re.compile(r"ST(\d+\.*\d*)-IN(\d+)\Z")
                steppattern = re.compile(r"ST(\d+\.*\d*)\Z")
                stepinitpattern = re.compile(r"ST(\d+\.*\d*)INIT\Z")
        
            for m in self.domain.meshes:
                
                nNodes = m.nnodes
                nElem = m.nElem
                
                pvdfile = param.WORKDIR + self.filename + '-mesh-' + str(m.UId)
                p = vtk.VtkGroup(pvdfile)
                
                for state in statekeys:
                    
                    vtufile = self.filename + '-mesh-' + str(m.UId) + '-' + state
                    
                    if re.match(iterpattern, state):
                        statetime = re.match(iterpattern, state).groups()
                        if len(statetime[1])==1:
                            statetime = list(statetime)
                            statetime[1] = '0'+statetime[1]
                        if len(statetime[2])==1:
                            statetime = list(statetime)
                            statetime[2] = '0'+statetime[2]                        
                        time = statetime[0] + '00' + statetime[1] + statetime[2]
                        level = 'iter'
                    elif re.match(incrpattern, state):
                        statetime = re.match(incrpattern, state).groups()
                        
                        if len(statetime[1])==1:
                            statetime = list(statetime)
                            statetime[1] = '0'+ statetime[1]
                        time = statetime[0] + '00' + statetime[1]
                        level = 'incr' 
                    elif re.match(steppattern, state):
                        statetime = re.match(steppattern, state).groups()
                        time = statetime[0]
                        level = 'step'
                    elif re.match(stepinitpattern, state):
                        statetime = re.match(stepinitpattern, state).groups()
                        time = statetime[0]
                        level = 'step'
                    else:
                        raise ValueError('state has bad value: '+str(state))
                    
                    if not level in levelfilter and levelfilter: continue
                    
                    outputfields, fieldscomp = param.getOutputFields(level)
                    
                    # build x, y and z of points
                    try:
                        
                        x, y, z = m.getNodesCooAtState(state, format='output', components=fieldscomp[outputfields.index('dof')], dtype='float64')
                    except:
                        x, y, z = m.getNodesCooAtState(format='output', dtype='float64')
                    
                    # build cells:
                    cell_connectivity = []
                    cell_offsets = []
                    cell_types = []
                    offset = 0
                    for i, elem in enumerate(m.elementsType):
                        nodes = [m.nodesMap['U2C'][nd] for nd in m.elementsNodes[i]]
                        if 'None' in nodes:
                            print 'bad element ', m.elementsMap['C2U'][i], ' of type ', elem, ': ', m.elementsNodes[i]
                        cell_connectivity.extend([nodes[k] for k in self.cells_connec_dict[elem]])
                        offset += len(nodes)
                        cell_offsets.append(offset)
                        cell_types.append(self.cells_type_dict[elem])
                     
                    cell_connectivity = np.array(cell_connectivity,dtype='int32')
                    cell_offsets = np.array(cell_offsets,dtype='int32')
                    cell_types = np.array(cell_types,dtype='uint8')
                    
                    # build point data: 'dof', 'force', ...
                    nodeData = {}
                    
                    for field in outputfields:
                        if field in self.node_fields:
                            comp = fieldscomp[outputfields.index(field)]
                            data = m.fields.getField(field, state)
                            for c in comp:
                                
                                nodeData[field + '-' + c] = data[:,int(c)].toarray().reshape(1,-1)[0,:]
                                
                    if not len(nodeData.keys()): nodeData = None
                    
                    # build cell data: 'stress', 'strain', 'fluxes', 
                    elemData = {}
                    
                    elemData['domain']=np.array([self.domain.getUId()]*nElem,dtype='int64')
                    
                    for field in outputfields:
                        if field in self.elem_fields:
                            comp = fieldscomp[outputfields.index(field)]
                            data = m.fields.getField(field, state)
                            for c in comp:
                                elemData[field + '-' + c] = data[:,int(c)]
                                
                    if not len(elemData.keys()): elemData = None  
                    
#                    onlypoints=True
#                    
#                    if onlypoints:
#                        print outputfields
#                        x, y, z = m.getNodesCooAtState(format='output', dtype='float64')
##                        npoints = 4
##                        x = np.array([0.,10.,20.,30.],dtype='float64')
##                        y = np.array([0.,10.,20.,30.],dtype='float64')
##                        z = np.array([0.,10.,20.,30.],dtype='float64')
#                        print x
#                        assert (x.size == y.size == z.size)
#                        npoints = x.size
#                        data=None
#                        # create some temporary arrays to write grid topology
##                        cell_offsets = np.arange(start = 1, stop = npoints + 1, dtype = 'int32')   # index of last node in each cell
##                        cell_connectivity = np.arange(npoints, dtype = 'int32')                    # each point is only connected to itself
##                        cell_types = np.empty(npoints, dtype = 'uint8') 
##                       
##                        cell_types[:] = vtk.VtkVertex.tid
#                        
#                        largeFileFlag = hl._requiresLargeVTKFileSize("VtkUnstructuredGrid", numPoints = npoints, numCells = npoints, pointData = data, cellData = None)
#                    
#                        w = vtk.VtkFile(param.WORKDIR + vtufile, vtk.VtkUnstructuredGrid, largeFileFlag)
#                        w.openGrid()
#                        w.openPiece(ncells = 40, npoints = npoints)
#                        
#                        w.openElement("Points")
#                        w.addData("points", (x,y,z))
#                        w.closeElement("Points")
#                        w.openElement("Cells")
#                        w.addData("connectivity", cell_connectivity)
#                        w.addData("offsets", cell_offsets)
#                        w.addData("types", cell_types)
#                        w.closeElement("Cells")
#                        
#                        hl._addDataToFile(w, cellData = None, pointData = data)
#                    
#                        w.closePiece()
#                        w.closeGrid()
#                        w.appendData( (x,y,z) )
#                        w.appendData(cell_connectivity).appendData(cell_offsets).appendData(cell_types)
#                    
#                        hl._appendDataToFile(w, cellData = None, pointData = data)
#                    
#                        w.save()
#                        
#                        p.addFile(vtufile + '.vtu',time)
#                    else:
#                    
#                        # build cell data: 'stress', 'strain', 'fluxes', 
#                        elemData = {}
#                        
#                        elemData['domain']=np.array([self.domain.getUId()]*nElem,dtype='int64')
#                        
#                        for field in outputfields:
#                            if field in self.element_fields:
#                                comp = fieldscomp[outputfields.index(field)]
#                                data = m.fields.getField(field, state)
#                                for c in comp:
#                                    elemData[field + '-' + c] = data[:,int(c)]
#                                    
#                        if not len(elemData.keys()): elemData = None  
#                          
#                        # Writing vtu file:
#                        largeFileFlag = hl._requiresLargeVTKFileSize("VtkUnstructuredGrid", numPoints = nNodes, numCells = nElem, pointData = nodeData, cellData = elemData)
#                        
#                        
#                        w =  vtk.VtkFile(param.WORKDIR + vtufile, ftype, largeFileFlag)
#                        w.openGrid()
#                        w.openPiece(npoints = nNodes, ncells = nElem)
#                           
#                        w.openElement("Points")
#                        w.addData("points", (x,y,z))
#                        w.closeElement("Points")
#                        
#                        w.openElement("Cells")
#                        w.addData("connectivity", cell_connectivity)
#                        w.addData("offsets", cell_offsets)
#                        w.addData("types", cell_types)
#                        w.closeElement("Cells")
#                        
#                        print w.xml.stream
#                        
#                        elemData = None
#                        pointData = None
#                        
#                        hl._addDataToFile(w, elemData, nodeData)
#                        w.closePiece()
#                        w.closeGrid()
#                        
#                        # Write grid data:
#                        w.appendData( (x,y,z) )
#                        w.appendData( cell_connectivity )
#                        w.appendData( cell_offsets )
#                        w.appendData( cell_types )
#                        
#                        # Write data:
#                        hl._appendDataToFile(w, elemData, nodeData)
#                        w.save()
                    
                    f = self.meshToVTK(param.WORKDIR + vtufile, x, y, z, cell_connectivity, cell_offsets, cell_types, nodeData, elemData)
                    p.addFile(f, time)
            
                p.save()

    def meshToVTK(self, path, x,y,z, connectivity, offsets, types, nodeData=None, elemData=None):
        
        ftype =vtk.VtkUnstructuredGrid
        
        assert (x.size == y.size == z.size)
        nNodes = x.size
        
        nElem = types.size
        
        largeFileFlag = hl._requiresLargeVTKFileSize("VtkUnstructuredGrid", numPoints = nNodes, numCells = nElem, pointData = nodeData, cellData = elemData)
                        
                        
        w =  vtk.VtkFile(path, ftype, largeFileFlag)
        w.openGrid()
        w.openPiece(npoints = nNodes, ncells = nElem)
           
        w.openElement("Points")
        w.addData("points", (x,y,z))
        w.closeElement("Points")
        
        w.openElement("Cells")
        w.addData("connectivity", connectivity)
        w.addData("offsets", offsets)
        w.addData("types", types)
        w.closeElement("Cells")
        
        hl._addDataToFile(w, elemData, nodeData)
        w.closePiece()
        w.closeGrid()
        
        # Write grid data:
        w.appendData( (x,y,z) )
        w.appendData( connectivity )
        w.appendData( offsets )
        w.appendData( types )
        
        # Write data:
        hl._appendDataToFile(w, elemData, nodeData)
        w.save()
        
        
        
        return w.getFileName()
        
        
#        gridToVTK(param.WORKDIR + "structured", x, y, z, cellData = {"pressure" : pressure}, pointData = {"temp" : temp})